package com.arithmetic.tree;

import java.util.*;

/**
 * 二叉树相关算法
 *
 * @author chenyichang
 * @since 2022/8/15
 */
public class TreeNodeArithmetic {
    /**
     * 层序遍历
     * 使用队列(先进先出)，只能在队尾添加元素
     */
    public static ArrayList<ArrayList<Integer>> levelOrder(TreeNode root) {
        ArrayList<ArrayList<Integer>> result = new ArrayList<>();
        if (root == null) {
            return result;
        }
        Queue<TreeNode> queue = new ArrayDeque<>();
        queue.add(root);

        //只要队列中不为空
        while (!queue.isEmpty()) {
            ArrayList<Integer> row = new ArrayList<>(queue.size());
            int size = queue.size();
            //1.取出队列的所有元素，并把元素值放入curList，
            for (int i = 0; i < size; i++) {
                TreeNode poll = queue.poll();
                row.add(poll.val);
                //2.并将元素的左右节点加入队列
                if (poll.left != null) {
                    queue.add(poll.left);
                }
                if (poll.right != null) {
                    queue.add(poll.right);
                }
            }
            result.add(row);
        }
        return result;
    }


    /**
     * 层序遍历，即一层一层输出
     * 用map方式，每次都需要排序，性能不行
     */
    public static ArrayList<ArrayList<Integer>> levelOrderWithMap(TreeNode root) {
        if (root == null) {
            return null;
        }
        Map<Integer, ArrayList<Integer>> result = new TreeMap<>(Integer::compareTo);
        levelOrder0WithMap(result, root, 1);
        return new ArrayList<>(result.values());
    }

    /**
     * @param level 当前层级
     */
    private static void levelOrder0WithMap(Map<Integer, ArrayList<Integer>> result, TreeNode root, int level) {
        if (root == null) {
            return;
        }
        result.computeIfAbsent(level - 1, k -> new ArrayList<>());

        //加入当前节点
        result.get(level - 1).add(root.val);

        levelOrder0WithMap(result, root.left, level + 1);
        levelOrder0WithMap(result, root.right, level + 1);
    }


    /**
     * 后续遍历：左->右->根
     */
    public static int[] postorderTraversal(TreeNode root) {
        if (root == null) {
            return new int[0];
        }
        List<Integer> result = new ArrayList<>();
        postorder(result, root);

        int[] ints = new int[result.size()];
        for (int i = 0; i < result.size(); i++) {
            ints[i] = result.get(i);
        }
        return ints;
    }

    private static void postorder(List<Integer> result, TreeNode root) {
        if (root == null) {
            return;
        }
        postorder(result, root.left);
        postorder(result, root.right);
        result.add(root.val);
    }


    /**
     * 中序遍历(左->根->右)
     */
    public static int[] inorderTraversal(TreeNode root) {
        // write code here
        if (root == null) {
            return new int[0];
        }
        List<Integer> result = new ArrayList<>();
        inorder(result, root);

        int[] ints = new int[result.size()];
        for (int i = 0; i < result.size(); i++) {
            ints[i] = result.get(i);
        }
        return ints;
    }

    private static void inorder(List<Integer> result, TreeNode root) {
        if (root == null) {
            return;
        }
        inorder(result, root.left);
        result.add(root.val);
        inorder(result, root.right);
    }


    /**
     * 二叉树前序遍历（根->左->右）
     * <p>
     * 1.输出当前节点
     * 2.输出当前节点的左节点
     * 3.输出当前节点的右节点
     */
    public static int[] preorderTraversal(TreeNode root) {
        if (root == null) {
            return new int[0];
        }
        List<Integer> result = new ArrayList();

        prePrint(result, root);

        int[] ints = new int[result.size()];
        for (int i = 0; i < result.size(); i++) {
            ints[i] = result.get(i);
        }
        return ints;
    }

    private static void prePrint(List<Integer> result, TreeNode curNode) {
        if (curNode == null) {
            return;
        }
        //加入当前节点
        result.add(curNode.val);
        //遍历左节点
        prePrint(result, curNode.left);
        //遍历右节点
        prePrint(result, curNode.right);
    }

}
